import subprocess
import glob
import os
import json
from loguru import logger
from typing import *
import sys
import platform

# Refactor V1
class Commander:
    # bash shell 工具链
    # readelf 工具链
    # grep 工具链
    # mkdir 工具链
    # dot 工具链， 生成图片
    tools = {} # 工具链配置

    @classmethod
    def get_tool(cls, key):
        """校验并返回工具链绝对路径
        """
        if key not in cls.tools:
            logger.error(f"{key} 未配置")
        return cls.tools[key]

    @classmethod
    def run(cls, cmd: str):
        """执行cmd命令
        """
        logger.debug(f"run cmd: {cmd}")
        result = subprocess.run(f'{cls.get_tool("bash")} -c "{cmd}"', shell=True, capture_output=True, text=True)
        logger.debug(f"cmd code: {result.returncode}")
        logger.debug("Standard output:", result.stdout)
        logger.debug("Standard error:", result.stderr)
        return [x for x in result.stdout.splitlines() if len(x) > 0 ]

    @classmethod
    def readelf(cls, lib, callback):
        """ 使用readelf 读取lib的依赖库

        lib: str 库的绝对路径
        """

        cmd = f"{cls.get_tool('readelf')} -a {lib} | {cls.get_tool('grep')} 'Shared library'"
        return callback(cls.run(cmd))
    
    @classmethod
    def otool(cls, lib, callback):
        """ 使用otool 读取lib的依赖库

        lib: str 库的绝对路径
        """
        cmd = f"{cls.get_tool('otool')} -L {lib} | {cls.get_tool('grep')} 'compatibility'"
        return callback(cls.run(cmd))

    @classmethod
    def readsymbol(cls, lib, callback):
        """ 读取 lib 的依赖库

        lib: str 库的绝对路径
        """
        if platform.system() == "Darwin":
            return cls.otool(lib, callback)
        else:
            return cls.readelf(lib, callback)

    @classmethod
    def mkdir(cls, path):
        """创建文件夹
        """
        cls.run(f"{cls.get_tool('mkdir')} -p {path}")

    @classmethod
    def dot(cls, src, dst):
        """生成图片
        """
        cls.run(f"{cls.get_tool('dot')} -Tsvg {src} -o {dst}")

class Vertex:
    def __init__(self, val: str):
        self.val = val

    def to_dot(self, transform=None) -> str:
        """生成dot格式节点字符串格式

        transform: 可选的转换函数
        """
        content = self.val
        if transform:
            content = transform(content)
        return f"\"{content}\""
    
    def ignore(self, filter) -> bool:
        """判断是否被忽略
        """
        return self.val in filter


    def __hash__(self):
        # 使用内置的 hash() 函数来构建哈希值
        return hash(self.val)
    
    def __eq__(self, other):
        if not isinstance(other, Vertex):
            return NotImplemented
        # 比较两个 Vertex 对象是否相等
        return self.val == other.val
    
    def __repr__(self):
        return f"Vertex({self.val})"

class Graph:
    def __init__(self, lib_depend_map, max_depth: int = 2):
        self.adj_list: Dict[Vertex, Vertex] = {} # 存储邻接表 Vertex -> Vertex[]
        self.rev_adj_list: Dict[Vertex, Vertex] = {} # 存储逆邻接表 Vertex -> Vertex[]
        # 构造邻接表和逆邻接表
        for (v0, depends) in lib_depend_map.items():
            self.adj_list[Vertex(v0)] = [Vertex(x) for x in depends]
        # 构建逆邻接表
        for (key, vals) in lib_depend_map.items():
            for val in vals:
                vertex = Vertex(val)
                if vertex not in self.rev_adj_list:
                    self.rev_adj_list[vertex] = []
                self.rev_adj_list[vertex].append(Vertex(key))
        # 最大深度
        self.max_depth = max_depth
        # visited 集合
        self.visited = set()
        logger.info(f"bulild graph success {len(self.adj_list)} {len(self.rev_adj_list)}!")


    def forward_visit(self, start: Vertex, action):
        """遍历从start开始的连通图
        """
        self.visited.clear()
        self.visit(start, self.visited, action, True)

    def backward_visit(self, end, action):
        """以end结束的连通图
        """
        self.visited.clear()
        self.visit(end, self.visited, action, False)

    def visit(self, ver: Vertex, visited, action, is_forawrd: bool, depth: int = 0):
        """通用遍历方法
        """
        # 深度超出最大深度
        if depth >= self.max_depth:
            return
        # 访问过直接退出
        if ver in visited:
            return
        if (is_forawrd and ver not in self.adj_list) or (not is_forawrd and ver not in self.rev_adj_list):
            # 前向遍历节点不在 邻接表里 或者 反向遍历 不在逆邻接表里
            return
        
        depend_libs = self.adj_list[ver] if is_forawrd else self.rev_adj_list[ver]
        
        # 相同接口的回调
        action(ver, depend_libs)

        visited.add(ver)

        # 递归遍历
        for lib in depend_libs:
            self.visit(lib, visited, action, is_forawrd, depth + 1)

class MyException(Exception):
    """这是一个自定义异常类"""
    def __init__(self, message):
        self.message = message
        super().__init__(self.message)

class DependAnalysis:
    """读取配置，获取相关上下文信息
    """
    def __init__(self, config="./config.json"):
        self.config = DependAnalysis.__read_json(config)
        # 中间结果
        self.lib_map = {} # str -> str 库名字到库路径的映射
        self.depend_map = {} # str -> str[] 库名字间的依赖关系
        # 使用 loguru 配置
        # logger.disable()
        self.init_common_config()


    def run(self):
        """根据commands配置执行命令
        """
        self.prepare()

        self.analysis()

    def prepare(self):
        """准备阶段: "init", "log", "collect", "buffer"
            , "load"
        """
        steps = self.get_cmd("prepare")
        for step in steps:
            if step == "init":
                self.init()
            elif step == "log" and os.path.exists(self.log_path):
                # 初始化log等级, 保证output目录已创建
                self.reset_log_level(self.log_level)
            elif step == "collect":
                # 收集库以及依赖
                self.collect()
            elif step == "buffer":
                # write lib_map & depend_map
                self.buffer()
            elif step == "load":
                # load from buffer
                self.load()
            else:
                # 不支持的命令
                logger.error(f"不支持的 prepare 命令, 请配置 {steps}")

    def analysis(self):
        """分析阶段
            forward: start -> xxx
            backward: xxx -> end
        """
        
        analysis_cfg = self.get_cmd("analysis")

        # 忽略库的集合
        ignore_libs = set()
        if "ignore" in analysis_cfg:
            ignore_libs.update(analysis_cfg["ignore"])
        
        threshold = 1000 # 默认1000KB，可配置
        if "threshold" in analysis_cfg:
            threshold = analysis_cfg["threshold"]
        max_depth = 2
        if "max_depth" in analysis_cfg:
            max_depth = analysis_cfg["max_depth"]

        logger.info(f"analysis: threshold: {threshold} kb, max_depth: {max_depth}")

        graph = Graph(self.depend_map, max_depth)

        # 标记列表
        mark_list = []

        def transform(content: str) -> str:
            ret = content
            if content in self.lib_map:
                lib_size = os.path.getsize(self.lib_map[content]) // 1024
                ret = f"{content}_{lib_size}_K"
                if lib_size > threshold: # 大于threshold标记
                    mark_list.append(f"\"{ret}\"") # mark
            return ret

        def depend(libs, is_forawrd=True):
            """依赖分析
            """
            separator = ", "
            for lib in libs:
                mark_list.clear()
                result = []
                vertex = Vertex(f"{lib}")
                result.append(f"{vertex.to_dot(transform)}[color=red]")
                def action(ver: Vertex, depend_libs: List[Vertex]):
                    depend_libs = [x for x in depend_libs if not x.ignore(ignore_libs)]
                    if not ver.ignore(ignore_libs) and len(depend_libs) > 0:
                        part0, part1 = f"{ver.to_dot(transform)}", f"{{{separator.join([x.to_dot(transform) for x in depend_libs])}}}"
                        if not is_forawrd:
                            part0, part1 = part1, part0
                        result.append(f"{part0} -> {part1};")
                if is_forawrd:
                    graph.forward_visit(vertex, action)
                else:
                    graph.backward_visit(vertex, action)
                for ml in mark_list:
                    result.append(f"{ml}[color=red]")
                suffix = "start" if is_forawrd else "end"
                self.out_dot(result, out_path=f"{lib}_{suffix}")
        
        # 前向，反向分析
        forwards = analysis_cfg["forward"] # 前向分析库的集合
        backwards = analysis_cfg["backward"]
        depend(forwards)
        depend(backwards, is_forawrd=False)

    def out_dot(self, lines, out_path="out"):
        """将lines写入dot文件, normal 一下文件名 . -> _
        """
        # write dot file
        out_path = out_path.replace(".", "_")
        with open(f"{self.output}/{out_path}.dot", "w") as out:
            out.write("digraph { \n")
            for line in lines:
                out.write(f"{line}\n")
            out.write("}\n")
        # export svg
        Commander.dot(f"{self.output}/{out_path}.dot", f"{self.output}/{out_path}.svg")
        logger.info(f"generate svg {self.output}/{out_path}.svg success")


    def init(self):
        """创建文件夹，配置自定义logger级别
        """
        # 创建 output & output/log
        Commander.mkdir(self.output)
        Commander.mkdir(self.log_path)

        logger.info("prepare init step success!")

    def init_common_config(self):
        """初始化一些全局配置
        """
        self.lib_map_path = self.get_config("lib_map_path")
        self.lib_depend_map_path = self.get_config("lib_depend_map_path")

        # 初始化工具链配置
        Commander.tools = self.get_config("tools")

        self.input_dirs = self.get_input("lib_paths")
        self.output = self.get_output("out")
        log = self.get_output("log") 
        self.log_level = self.get_output("log_level") # DEBUG, INFO, ERROR
        self.log_path = f"{self.output}/{log}"

        self.log_handlers_ids = [] # 初始化log handlers
        logger.remove() # 清空默认的处理器
        

    def reset_log_level(self, level):
        """ 动态设置 loguru 的日志级别
        """
        # 移除所有已有的处理器
        for handler_id in self.log_handlers_ids:
            logger.remove(handler_id)
        
        # 清空 handlers_ids 列表，准备添加新的处理器
        self.log_handlers_ids.clear()
        
        # 根据新的日志级别重新添加处理器
        self.log_handlers_ids.append(logger.add(sys.stderr, level=level))
        self.log_handlers_ids.append(logger.add(f"{self.log_path}/{level.lower()}.log", level=level))

        logger.info("prepare log step success!")
 
    def collect(self):
        """收集所有库的名字
        """
        for dir in self.input_dirs:
            self.build_lib_map(dir)
        # 构建依赖图
        self.build_depend_map()

        logger.info("prepare collect step success!")


    def buffer(self):
        """使用json持久化中间结果: lib_map & depend_map
        """
        self.check_buffer_config()
        with open(self.lib_map_path, "w") as jsonfile:
            json.dump(self.lib_map, jsonfile)
        with open(self.lib_depend_map_path, "w") as jsonfile:
            json.dump(self.depend_map, jsonfile)

        logger.info("prepare buffer step success!")


    def load(self):
        """使用json从缓存中加载中间结果: lib_map & depend_map
        """
        # check 配置
        self.check_buffer_config()
        with open(self.lib_map_path, "r") as jsonfile:
            self.lib_map = json.load(jsonfile)
        with open(self.lib_depend_map_path, "r") as jsonfile:
            self.depend_map = json.load(jsonfile)

        logger.info(f"prepare load step success {len(self.lib_map)} {len(self.depend_map)}!")

    def check_buffer_config(self):
        # check 配置
        if self.lib_depend_map_path == "" or self.lib_map_path == "":
            logger.error(f"buffer path 没有配置")

    def build_lib_map(self, dir: str):
        """获取dir下的所有库文件(递归)
        """
        # 获取当前目录及其子目录下的所有 .so 文件
        # pattern = f'{dir}/**/*.so'
        # 知识点: 这里的dir 在windwos上必须使用 C: 开头的格式
        if platform.system() == "Darwin":
            pattern = os.path.join(dir, '**', '*.dylib')
        else:
            pattern = os.path.join(dir, '**', '*.so')
        tmp_libs = glob.glob(pattern, recursive=True)
        logger.debug(f"get libs: {len(tmp_libs)}")
        for lib in tmp_libs:
            lib = lib.replace("\\", "/") # normal 路径中的分隔符 /
            libname = lib.split("/")[-1] # 获取标准路径的文件名(带扩展)
            if libname not in self.lib_map:
                self.lib_map[libname] = lib # 库名 -> 库绝对路径(normal)
            else:
                logger.warning(f"{libname} is duplicated")
        logger.info(f"build lib_map success: {len(self.lib_map)}")

    def build_depend_map(self):
        """根据lib_map 构建依赖图
        """
        for (libname, libpath) in self.lib_map.items():
            # lines 是命令的日志输出 -> 依赖库列表[]
            callback = lambda lines: [] if len(lines) == 0 else [self.get_depend_lib(x) for x in lines]

            depend_libs = Commander.readsymbol(libpath, callback)

            depend_libs = [x for x in depend_libs if x != libname]

            logger.debug(f"{libname} depends: {len(depend_libs)}")
            
            self.depend_map[libname] = depend_libs
        logger.info(f"build depend_map success: {len(self.depend_map)}")
            
    def get_depend_lib(self, line):
        """根据日志行获取依赖库名
        
        line: 0x0000000000000001 (NEEDED)             Shared library: [libdevmapper.so.1.02.1]
        
        line: @rpath/libcangjie-std-ast.dylib (compatibility version 0.0.0, current version 0.0.0)
        return
            libdevmapper.so.1.02.1
        """
        if "Shared library" in line:
            lib_column = line.split()[-1]
            return lib_column[1:len(lib_column)-1]
        elif "compatibility" in line:
            # mac
            lib_column = line.split()[0]
            return lib_column.split("/")[-1]

        return ""
    
    def get_cmd(self, key):
        return self.get_config("commands", key)
    
    def get_input(self, key):
        return self.get_config("input", key)

    def get_output(self, key):
        return self.get_config("output", key)

    @logger.catch
    def get_config(self, key0, key1=None) -> str:
        """获取配置

        key1 可选
        """
        if key0 not in self.config:
            raise MyException(f"{key0} 未配置")

        result = self.config[key0]
        if key1 is None:
            return result
        
        if key1 not in self.config[key0]:
            raise MyException(f"{key1} 未配置 in {key0}")
            
        return result[key1]

    @staticmethod
    def __read_json(config):
        with open(config, 'r') as jsonfile:
            return json.load(jsonfile)
    

def run():
    # reactor v1 test
    dt = DependAnalysis()
    dt.run()

# D:/dev/py310/python.exe
if __name__ == "__main__":
    run()
